import os
import sqlite3

from dotenv import load_dotenv
from langchain_community.utilities import SQLDatabase
from langchain_community.agent_toolkits import SQLDatabaseToolkit
from langchain_core.messages import SystemMessage
from langchain_core.messages import HumanMessage
from langgraph.prebuilt import create_react_agent

from lc_frame.lc_models.api import llm_api
from lc_frame.utils.log import logger

current_path = os.path.abspath(__file__)
print(current_path)
root_path = os.path.dirname(os.path.dirname(current_path))

# 加载key
dotenv_path = os.path.join(root_path, "lc_models", "api", "qwen.txt")
load_dotenv(dotenv_path=dotenv_path)

# 连接本地数据库
#relational_db_path = os.path.join(root_path, "datas", "relational_db", "博金杯比赛数据.db")
relational_db_path="/gemini/code/arua_completion_project/dataset/博金杯比赛数据.db"
db = SQLDatabase.from_uri(f'sqlite:///{relational_db_path}')
Trade_date=[]

def get_trading_date(current_trading_date):
    """
       获取上一交易日
       入参：当前交易日
       反参：上一交易日   
    """
    last_trading_date = ""
    if len(Trade_date)>0:
        pass
    else:
        conn = sqlite3.connect(relational_db_path)
        cursor = conn.cursor()
        searchsql="""
        SELECT distinct "交易日" from 'A股票日行情表' order by "交易日" asc
        """
        
        cursor.execute(searchsql)
        rows = cursor.fetchall()
        for row in rows:
            Trade_date.append(row[0])
        print(Trade_date)
        cursor.close()
        conn.close()
    # 找到目标元素的索引
    try:
        index = Trade_date.index(current_trading_date)
        # 检查索引是否大于 0
        if index > 0:
            last_trading_date = Trade_date[index - 1]
        else:
            print("这是列表的第一个元素，没有前一个元素")
    except ValueError:
        print(f"元素 {current_trading_date} 不在列表中")
    return last_trading_date


def lc_agent(question):
    # 接入千问大模型
    chat, llm, embeddings = llm_api.get_qwen()

    # 调用langchain的数据库工具
    toolkit = SQLDatabaseToolkit(db=db, llm=chat)
    tools = toolkit.get_tools()
    tools.append(get_trading_date)

    # 提示词数据库模版
    # api https://api.python.langchain.com/en/latest/messages/langchain_core.messages.system.SystemMessage.html
    SQL_PREFIX = """You are an agent designed to interact with a SQL database.
    Given an input question, create a syntactically correct SQLite query to run, then look at the results of the query and return the answer.
    Unless the user specifies a specific number of examples they wish to obtain, always limit your query to at most 3 results.
    You can order the results by a relevant column to return the most interesting examples in the database.
    Never query for all the columns from a specific table, only ask for the relevant columns given the question.
    You have access to tools for interacting with the database.
    Only use the below tools. Only use the information returned by the below tools to construct your final answer.
    You MUST double check your query before executing it. If you get an error while executing a query, rewrite the query and try again.
    DO NOT make any DML statements (INSERT, UPDATE, DELETE, DROP etc.) to the database.
    To start you should ALWAYS look at the tables in the database to see what you can query.
    Do NOT skip this step.
    Then you should query the schema of the most relevant tables.
    请注意：
    1、如果没有获取到信息，或者查询时间超过10秒，返回没有查询到结果即可。
    2、如果生成的SQL语句中，字段带有英文括号()，请使用双引号包裹起来，例如：收盘价(元) 双引号包裹为 "收盘价(元)"。
    3、如果查询过程中SQL语句有语法错误，减少查询量，总体查询次数应控制在5次以内。
    4、在联表查询时，每次使用'A股公司行业划分表'时都需要使用'A股公司行业划分表'表中的'交易日期'字段进行查询，而不是仅使用表'A股票日行情表'或表'港股票日行情表'中的'交易日'字段进行查询。
    5、查询"涨跌幅"、"日收益率"时，优先使用系统表中的结果。
    6、"涨跌幅">=0时为"涨幅"，<0时为"跌幅",查询最大或最小涨跌幅时,要取绝对值。
    7、查询完毕后请及时关闭数据库查询链接。
    8、如果问题不涉及行业信息，请不要关联"A股公司行业划分表"进行查询。
    9、所有日期请用以下格式，如：'20190224'。
    10、开盘价较上一交易日高指同一只股票当日的"开盘价(元)"比上一交易日的"最高价(元)"大
    """
    system_message = SystemMessage(content=SQL_PREFIX)

    # 构建代理
    agent_executor = create_react_agent(chat, tools=tools, state_modifier=system_message)
    # 流式查询
    events = agent_executor.stream(
        {"messages": [HumanMessage(content=question)]},
        stream_mode="values"
    )
    result_list = []
    for event in events:
        result_list.append(event["messages"][-1].content)
        logger.info(event["messages"][-1].pretty_print())
    # 返回最终结果
    final_result = event["messages"][-1].content if result_list else None
    logger.info(f'最终答案是：{final_result}')
    return final_result


if __name__ == "__main__":
    question = "基金类型为货币型的基金有哪些？"
    lc_agent(question)
